#ifndef M_PT
#define M_PT
#include<omp.h>
#include"raytracer.h"
#include"../media/homogeneous.h"
using namespace std;

// Path Tracing
class PT :public RayTracer
{
private:
	Spectrum PathTrace(int pix, Ray ray, RandGenerator *randGen);
	inline Float clamp(Float x) { return x<0 ? 0 : x>1 ? 1 : x; }
	inline int toInt(Float x) { return int(pow(clamp(x), 1 / 2.2) * 255 + .5); }

public:
	PT(bool enableMedia) :RayTracer(RT_PT,enableMedia) {}
	~PT(){}
	void render(string outputname,string config);
};

Spectrum PT::PathTrace(int pix, Ray ray, RandGenerator *randGen)
{
	Float etaScale = 1;
	Spectrum beta(1.f);
	bool isSpec = false;
	Spectrum luminance(0);
	Ray current_ray = ray;
	int loopcnt = 0;

	for (int depth = 0; depth < 8; ++depth)
	{
		Float t = INF;
		//Medium *nextMedium=current_ray.medium;
		SurfaceInteraction si;
		bool intersect = scene->nearest_intersect(current_ray, &t, &si);
		Object* pObj = static_cast<Object*>(si.pObj);
		bool sampleMedia=false;
		P3 dir=current_ray.d, its_pt=current_ray.o;

		if (enableMedia&&current_ray.medium)
		{
			Float sampled_t;
			beta *= current_ray.medium->sampleMediumInteraction(current_ray, t, &sampled_t, randGen);
			if (beta.IsBlack()) break;
			sampleMedia = sampled_t < t;
			its_pt = current_ray.get(sampled_t);
		}
		if (sampleMedia)
		{
			
			if (intersect)
			{
				//get bsdf
				Material* material = pObj->material;
				shared_ptr<BSDF> bsdf;
				if (material)
					bsdf = material->createBSDF(&si, TransportMode::Radiance);
				//choose light
				Float u0, lightPdf;
				u0 = randGen->genRand();
				int lightId = min((int)(u0*scene->lights.size()), scene->lights.size() - 1);
				lightPdf = 1.f / (scene->lights.size());
				Light* light = scene->lights[lightId];
				Ray isect_ray = { its_pt,current_ray.d,current_ray.medium };
				luminance += beta * sampleDirectLighting(&si, bsdf, isect_ray, light, randGen, true) / lightPdf;
			}

			Float pdf;
			Float rands[] = { randGen->genRand(),randGen->genRand() };
			dir = current_ray.medium->sample_inScatterDir(its_pt, -current_ray.d, &pdf, rands);
			isSpec = false;
		}
		else
		{
			if (depth == 0 || isSpec)
			{
				//accumulate environment light
				if (!intersect)
				{
					for (auto const &light : scene->lights)
					{
						luminance += beta * light->lightEnv(current_ray);
					}
					break;
				}
				else
				{
					//to be refined: accumulate emission from object
					auto Le= beta * pObj->lightOut(&si, -current_ray.d);
					luminance += Le;
					if (!Le.IsBlack()) break;  // to be checked
				}
			}
			if (!intersect || !pObj ||loopcnt>30)
			{
				break;
			}
			MediumInterface mi = pObj->mi;
			its_pt = current_ray.get(t);
			Material* material = pObj->material;
			if (!material)
			{
				current_ray.o = current_ray.get(t+eps);
				current_ray.medium = mi.getMedium(si.n,current_ray.d);
				--depth;
				++loopcnt;
				continue;
			}
			shared_ptr<BSDF> bsdf = material->createBSDF(&si, TransportMode::Radiance);
			if (bsdf->getNumAllComponents() == 0) {
				current_ray.o = current_ray.get(t + eps);
				current_ray.medium = mi.getMedium(si.n, current_ray.d);
				--depth;
				++loopcnt;
				continue;
			}
			if (bsdf->getNumComponent(BxDFType(BSDF_ALL & ~BSDF_SPECULAR))>0)
			{
				// accumulate direct lighting
				Spectrum ld;


				////choose one light to sample based on light power distribution
				//Float u0 = randGen->genRand(), lightPdf;
				//int lightId = scene->lightPowerDistribution->sampleDiscrete(u0, &lightPdf);
				//Light* light = scene->lights[lightId];

				Float u0 = randGen->genRand();
				int lightId = fmin((int)(u0 * scene->lights.size()), scene->lights.size() - 1);
				Float lightPdf = 1.f / (scene->lights.size());
				Light* light = scene->lights[lightId];

				Ray isect_ray = { si.its_pt,current_ray.d,current_ray.medium };
				luminance += beta * sampleDirectLighting(&si, bsdf, isect_ray, light, randGen, false) / lightPdf;

				//P3 wi;
				//VisibilityTester vis;
				//Float pdf1;
				//Float rands1[] = { randGen->genRand(),randGen->genRand() };
				//Spectrum light_emission = light->sampleLightPt(its_pt, &wi, rands1, &pdf1, &vis) / lightPdf;
				//if (pdf1 > 0 && !light_emission.IsBlack())
				//	if (vis.unoccluded(scene, id))
				//	{
				//		ld += beta * (light_emission*abs(wi.dot(n)) / pdf1)*bsdf->f(-wi, -current_ray.d);
				//	}
				//luminance += ld;

			}
			Float pdf;

			Float rands[] = { randGen->genRand(),randGen->genRand() };
			BxDFType sampledType;
			Spectrum f = bsdf->ImportanceSample_f(-current_ray.d, &dir, rands, &pdf, BxDFType(BSDF_ALL), &sampledType);



			if (f.IsBlack() || pdf == 0.f)
			{
				//fprintf(stderr, "are you good malaysia\n");
				break;
			}
			isSpec = (sampledType&BSDF_SPECULAR);

			if ((sampledType & BSDF_SPECULAR) && (sampledType & BSDF_TRANSMISSION))
			{
				Float eta = bsdf->eta;
				// Update the term that tracks radiance scaling for refraction
				// depending on whether the ray is entering or leaving the
				// medium.
				etaScale *= (si.n.dot(-current_ray.d) > 0) ? (eta * eta) : 1.0 / (eta * eta);
			}

			//if ((sampledType&BSDF_SPECULAR) == 0) //not specular
			{
				f = f * abs(si.n.dot(dir)) / pdf;
			}

			beta *= f;
		}

		// Russian roulette
		Spectrum rrBeta = beta * etaScale;
		Float maxcomp = rrBeta.MaxComponentValue();
		if (maxcomp < .25f&&depth>3)
		{
			Float q = max(.05f, 1.f - maxcomp);
			if (randGen->genRand() < q) break;
			beta /= (1 - q);
		}
		MediumInterface mi = pObj->mi;
		current_ray.o = its_pt + dir * eps;
		current_ray.d = dir;
		current_ray.medium = mi.getMedium(si.n, dir);
		//if(n.dot(dir)<0)
		//	current_ray.medium = nextMedium;
	}
	if (isnan(luminance.MaxComponentValue())) return 0.;
	return luminance;
}

void PT::render(string outputname, string config)
{
	int w = cam->w, h = cam->h;
	int samps; // # samples 
	if (config == "")
	{
		samps = 16;
	}
	else
	{
		sscanf(config.c_str(), "%d", &samps);
	}

	//Spectrum sigma_a, sigma_s;
	//GetMediumScatteringProperties("Haze", &sigma_a, &sigma_s);
	//PhaseFunction *hg = new HenyeyGreenstein(0.7);
	//Medium *shampoo = new HomogeneousMedium(sigma_a, sigma_s, hg);

	//Spectrum r;
	P3 *c = new P3[w*h];
#pragma omp parallel for schedule(dynamic, 1)      // OpenMP 
	for (int y = 0; y<h; y++) {
		fprintf(stderr, "\rRendering (%d spp) %5.2f%%", samps * 4, 100.*y / (h - 1));
		for (int x = 0; x < w; x++)
		{
			int index = w * (h - 1 - y) + (w - 1 - x);
			//if (x != 119 || y != 162) continue;
			Spectrum r;
			for (int sy = 0; sy < 2; sy++) 
			{
				// 2x2 subpixel rows 
				for (int sx = 0; sx < 2; sx++)
				{        // 2x2 subpixel cols 
					for (int s = 0; s < samps; s++)
					{
						//long long seed = sx*s + 2*sy*s + (x*x+y*y);
						//long long seed = sx + 2 * sy + (x*x + y * y)*s+primes[s];
						long long seed = (s + 1) * (sx + 1) * (sy + 1) * (w * y + x);
						//srand(seed);
						DefaultRandGenerator haltonGen(seed);
						Float r1 = 2 * haltonGen.genRand(), dx = r1 < 1 ? sqrtf(r1) - 1 : 1 - sqrtf(2 - r1);
						Float r2 = 2 * haltonGen.genRand(), dy = r2 < 1 ? sqrtf(r2) - 1 : 1 - sqrtf(2 - r2);
						Float tx = (Float)(x + (sx + dx + .5f) / 2) / w, ty = (Float)(y + (sy + .5f + dy) / 2) / h;
						Ray ray = cam->gen_ray(tx, ty, NULL);
						//ray.medium = shampoo;
						r = r + PathTrace(index, ray, &haltonGen) * (1.f / samps);
					}
				}
			}
			P3 rgb;
			r.ToRGB((Float*)&rgb);
#pragma omp critical 
			{
			c[index] = c[index] + rgb*0.25f;
			}
		}
	}


	char filename[100];
	sprintf(filename, "out/%s.bmp", outputname.c_str());
	save_to_bmp(filename, w, h, &c[0].x, 2.2);

}

#endif // !M_PT

